#include "postprocess/processor/pedestrian_processor.h"

PedestrianProcessor::PedestrianProcessor() noexcept
    : BaseProcessor(), matcher_(nullptr), is_init_(false) {}

PedestrianProcessor::~PedestrianProcessor() {}

void PedestrianProcessor::Init(PostProcessorSettings const &settings) noexcept {
  std::lock_guard<std::mutex> lock(mutex_);
  if (is_init_) return;
  settings_ = settings;
  // do init
  is_init_ = true;
  pre_timestamp_ = 0;
  track_initialization_flag_ = false;
  pre_pose_.matrix() = Eigen::Matrix4d::Identity();
  cur_pose_.matrix() = Eigen::Matrix4d::Identity();
  processor_flag_ = ProcessorClass::Pedestrian;
  matcher_.reset(new Mathcer(processor_flag_, settings_));
  // add class flag
  const std::vector<uint32_t> vec{5, 6, 7, 8};
  SetClassFlag(vec);
}

bool PedestrianProcessor::Run(
    autodrive::perception::PerceptionResult &perceptionResult,
    Eigen::Affine3d &pos, std::string &err) noexcept(false) {
  if (!is_init_) {
    err = "pedestrian post processor module must do init  at first!";
    return false;
  }

  if (class_vec_.empty()) {
    err = "pedestrian post processor module must set class flag at first!";
    return false;
  }
  cur_pose_ = pos;
  // for first frame ,tracker need to be initialized
  if (!track_initialization_flag_) {
    std::vector<std::int32_t> objects_index;
    std::vector<TrackedObjectPtr> tracked_objects;
    construct_tracked_objects(
        perceptionResult.dl_boxes, &tracked_objects,
        perceptionResult.timestamp);  // 依据类别，从dl_boxes生成tracked_objects
    range(perceptionResult.dl_boxes, objects_index);  // 生成0~(n-1)的序列
    create_new_tracks(tracked_objects, objects_index, object_tracks_,
                      settings_);  // 生成track，将lidar_object转至track_object_
    // collect_tracked_results(perceptionResult, object_tracks_);
    track_initialization_flag_ = true;
    pre_timestamp_ = perceptionResult.timestamp;
    pre_pose_ = cur_pose_;
    return true;
  } else {
    std::uint64_t timestamp{perceptionResult.timestamp};
    double time_diff{static_cast<double>(timestamp - pre_timestamp_) /
                     1000000.0f};
    pre_timestamp_ = timestamp;
    Eigen::Matrix4f transformed_matrix =
        (cur_pose_.matrix().inverse() * pre_pose_.matrix()).cast<float>();
    // A. construct objects for tracking
    std::vector<TrackedObjectPtr> tracked_objects;
    construct_tracked_objects(perceptionResult.dl_boxes, &tracked_objects,
                              perceptionResult.timestamp);

    // B. prediction
    std::vector<Eigen::VectorXf> tracks_predict;
    compute_tracks_predict(tracks_predict, time_diff, object_tracks_);

    // C. match objects to tracks
    std::vector<std::pair<int, int>> assignments;
    std::vector<int> unassigned_tracks;
    std::vector<int> unassigned_objects;
    std::vector<ObjectTrackPtr> &tracks = object_tracks_.GetTracks();
    if (matcher_ != nullptr) {
      matcher_->Match(&tracked_objects, tracks, tracks_predict, &assignments,
                      &unassigned_tracks, &unassigned_objects,
                      transformed_matrix);  // match object and track
    }

    // D. Update Tracks
    // D.1 update tracks with associated objects
    update_assigned_tracks(&tracked_objects, assignments, time_diff,
                           object_tracks_);

    // D.2 update tracks without associated objects
    update_unassigned_tracks(tracks_predict, unassigned_tracks, time_diff,
                             object_tracks_);

    // D.3 delete track
    delete_lost_tracks(object_tracks_);

    // D.4 create new tracks for objects without associated tracks
    create_new_tracks(tracked_objects, unassigned_objects, object_tracks_,
                      settings_);

    // D.5 set position
    pre_pose_ = cur_pose_;
    return true;
  }
}
